#include "HookLibrary.h"
#include "HookInfoFile.h"
#include "FileWatcher.h"

#define BufferSize (0x1000)

static volatile FileWatcherStatus WatcherStatus = FileWatcherStopped;

static volatile LPBYTE  Buffer         = NULL;
static volatile LPWSTR  FileName       = NULL;
static volatile DWORD   FileNameLength = 0;
static volatile HANDLE  Directory      = INVALID_HANDLE_VALUE;
static volatile HMODULE PinnedModule   = NULL;
static volatile HANDLE  WatcherThread  = NULL;
static volatile HWND    OwnerWindow    = NULL;

//
// Callback APC
//
VOID CALLBACK Callback(DWORD ErrorCode, DWORD NumberOfBytesTransferred, LPOVERLAPPED Overlapped)
{
	PFILE_NOTIFY_INFORMATION Info;

	DWORD Len;
	PBYTE NextEntry = Buffer;

	UNREFERENCED_PARAMETER(Overlapped);

	if ((ErrorCode != 0) || (NumberOfBytesTransferred == 0))
	{
		return;
	}

	do
	{
		Info = (PFILE_NOTIFY_INFORMATION)(NextEntry);

		if ((Info->Action == FILE_ACTION_ADDED) ||
			(Info->Action == FILE_ACTION_RENAMED_NEW_NAME))
		{
			Len = ((Info->FileNameLength) / sizeof(WCHAR));

			if (Len == FileNameLength)
			{
				if (StrCmpNIW(Info->FileName, FileName, Len) == 0)
				{
					OnHookInfoChanged(OwnerWindow);
				}
			}
		}

		NextEntry += Info->NextEntryOffset;
	}
	while (Info->NextEntryOffset != 0);
}

//
// Directory notifier
//
HANDLE WINAPI OpenHookInfoDirectory()
{
	WCHAR Path[MAX_PATH];

	if (!GetHookInfoDirectory(Path))
	{
		return INVALID_HANDLE_VALUE;
	}

	return CreateFileW(Path,
		               FILE_LIST_DIRECTORY,
					   FILE_SHARE_READ | FILE_SHARE_WRITE,
					   NULL,
					   OPEN_EXISTING,
					   FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OVERLAPPED,
					   NULL);
}

BOOL WINAPI ReadHookInfoDirectoryChanges(LPOVERLAPPED Overlapped)
{
	return ReadDirectoryChangesW(Directory,
						         Buffer,
						         BufferSize,
						         FALSE,
						         FILE_NOTIFY_CHANGE_FILE_NAME,
						         NULL,
						         Overlapped,
						         Callback);
}


//
// Watcher thread
//
DWORD WINAPI WatcherThreadProc(LPVOID Parameter)
{
	OVERLAPPED Overlapped;

	HANDLE OwnerThread = (HANDLE)(Parameter);
	DWORD  Result      = WAIT_FAILED;

	do
	{
		if (!ReadHookInfoDirectoryChanges(&Overlapped))
		{
			break;
		}

		Result = WaitForSingleObjectEx(OwnerThread, INFINITE, TRUE);
	}
	while (Result == WAIT_IO_COMPLETION);

	StopFileWatcher();
	CloseHandle(OwnerThread);
	FreeLibraryAndExitThread(PinnedModule, 0);
}
//
// Start/Stop
//
HMODULE WINAPI PinLibraryInProcess(HMODULE Module)
{
	HMODULE Result;

	if (GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, (LPCWSTR)(Module), &Result))
	{
		return Result;
	}

	return NULL;
}

BOOL WINAPI StartFileWatcherThread(HANDLE hWndOwner)
{
	HANDLE OwnerThread = NULL;
	DWORD  OwnerThreadId = GetWindowThreadProcessId(hWndOwner, NULL);

	if (OwnerThreadId != 0)
	{
		OwnerThread = OpenThread(SYNCHRONIZE, FALSE, OwnerThreadId);
		
		if (OwnerThread != NULL)
		{
			WatcherThread = CreateThread(NULL,
										 0x10000,
										 WatcherThreadProc,
										 OwnerThread,
										 CREATE_SUSPENDED,
										 NULL);

			if (WatcherThread != NULL)
			{
				SetThreadPriority(WatcherThread, THREAD_PRIORITY_IDLE);

				PinnedModule = PinLibraryInProcess(g_hModule);
				
				if (PinnedModule != NULL)
				{
					OwnerWindow = hWndOwner;

					if (ResumeThread(WatcherThread) != (DWORD)(-1))
					{
						return TRUE;
					}

					FreeLibrary(PinnedModule);
					PinnedModule = NULL;
					OwnerWindow = NULL;
				}

				TerminateThread(WatcherThread, 0);
				CloseHandle(WatcherThread);
				WatcherThread = NULL;
			}

			CloseHandle(OwnerThread);
			OwnerThread = NULL;
		}
	}

	return FALSE;
}

BOOL WINAPI StartFileWatcher(HWND hWndOwner)
{
	BOOL Succeeded = TRUE;

	Buffer = BGHeapAlloc(BufferSize);
	if (Buffer == NULL)
	{
		Succeeded = FALSE;
		goto Exit;
	}

	FileName = BGHeapAllocStr(MAX_PATH);
	if (FileName == NULL)
	{
		Succeeded = FALSE;
		goto Exit;
	}

	FileNameLength = GetHookInfoFileName(hWndOwner, FileName, MAX_PATH);
	if (FileNameLength == 0)
	{
		Succeeded = FALSE;
		goto Exit;
	}

	Directory = OpenHookInfoDirectory();
	if (Directory == INVALID_HANDLE_VALUE)
	{
		Succeeded = FALSE;
		goto Exit;
	}

	Succeeded = StartFileWatcherThread(hWndOwner);

Exit:

	if (!Succeeded)
	{
		StopFileWatcher();
	}

	return Succeeded;
}

BOOL WINAPI StopFileWatcher()
{
	BOOL Result = TRUE;

	if (Directory != INVALID_HANDLE_VALUE)
	{
		Result   &= CloseHandle(Directory);
		Directory = INVALID_HANDLE_VALUE;
	}

	if (FileName != NULL)
	{
		Result  &= BGHeapFree(FileName);
		FileName = NULL;
	}

	if (Buffer != NULL)
	{
		Result &= BGHeapFree(Buffer);
		Buffer  = NULL;
	}
	
	return Result;
}
VOID WINAPI EnsureFileWatcher(HWND hWnd)
{
	if (WatcherStatus == FileWatcherStopped)
	{
		if (StartFileWatcher(hWnd))
		{
			WatcherStatus = FileWatcherRunning;
			OnHookInfoChanged(hWnd);
		}
		else
		{
			WatcherStatus = FileWatcherFailed;
		}
	}
}